import xml.etree.ElementTree as ET
import os
import glob

'''
写入txt格式为

xml对应图片绝对路径   含物体个数   第一个物体对应bbox对应4个坐标       第一个物体对应bbox对应物体名称
0001.jpg               1           Xmin Ymin Xmax Ymax    破洞
'''


def create(root_path):
    # 新建一个名为all_info.csv文件，准备写入数据
    root = os.path.abspath(root_path) + '/'
    txt_file = open(root + "all_info.csv", mode='w', encoding='utf-8')

    # 拿到该文件夹下所有jpg文件
    jpg_paths = glob.glob(os.path.join(root, '*.xml'))
    jpg_paths.sort()

    # 遍历所有jpg, xml文件，将文件写入txt中
    for jpg_path in jpg_paths:
        # str.replace('a','b') 将字符串中的a替换为b
        path_split = jpg_path.split('/')
        jpg_name = path_split[-1]
        xml_name = jpg_path
        # txt 写入图像名字
        txt_file.write(jpg_name + ',')
        # results保存的是bbox物体的名字与对应4个坐标
        results = __parse_rec(xml_name)
        # num_obj 一张图中的物体总数
        num_obj = len(results)
        # txt写入  一张图中的物体总数
        txt_file.write(str(num_obj) + ',')
        # 遍历一张图片中的所有物体
        for result in results:
            name = result['name']
            bbox = result['bbox']
            # txt写入   bbox的坐标 以及  每个物体对应的类
            txt_file.write(
                str(bbox[0]) + ',' + str(bbox[1]) + ',' + str(bbox[2]) + ',' + str(bbox[3]) + ',' + name + ',')
        # 读取完一张图片后换行
        txt_file.write('\n')
    # 关闭txt
    txt_file.close()
    print(root + "all_info.csv", '生成成功')


def __parse_rec(filename):
    """
    解析一个 xml 中所有瑕疵位置
    """

    # 存储一张图片中的所有物体
    objects = []
    try:
        tree = ET.parse(filename)
        # 遍历一张图中的所有物体
        for obj in tree.findall('object'):
            obj_struct = {}
            # 读取物体名称
            obj_struct['name'] = obj.find('name').text
            bbox = obj.find('bndbox')
            # 从原图左上角开始为原点，向右为x轴，向下为y轴。左上角（xmin，ymin）和右下角(xmax,ymax)
            obj_struct['bbox'] = [int(float(bbox.find('xmin').text)),
                                  int(float(bbox.find('ymin').text)),
                                  int(float(bbox.find('xmax').text)),
                                  int(float(bbox.find('ymax').text))]
            objects.append(obj_struct)
    except FileNotFoundError:
        objects.append({'name': '正常', 'bbox': [0, 0, 0, 0]})

    return objects



STATISTICS = {
    '正常': [0, 0],
    '扎洞': [0, 0],
    '毛斑': [0, 0],
    '擦洞': [0, 0],
    '毛洞': [0, 0],
    '织稀': [0, 0],
    '吊经': [0, 0],
    '缺经': [0, 0],
    '跳花': [0, 0],
    '油/污渍': [0, 0],
    '其他': [0, 0],
}

# 新建一个名为all_info.csv文件，准备写入数据
root = os.path.abspath("../PascalVOC/val/Annotations") + '/'
# 拿到该文件夹下所有jpg文件
jpg_paths = glob.glob(os.path.join(root, '*.xml'))
jpg_paths.sort()

# 遍历所有jpg, xml文件，将文件写入txt中
for jpg_path in jpg_paths:
    # str.replace('a','b') 将字符串中的a替换为b
    path_split = jpg_path.split('/')
    jpg_name = path_split[-1]
    xml_name = jpg_path
    # results保存的是bbox物体的名字与对应4个坐标
    results = __parse_rec(xml_name)
    for result in results:
        for key in STATISTICS:
            if str(result).__contains__(key):
                STATISTICS[key][1] += 1
        if str(result).__contains__("渍"):
            STATISTICS["油/污渍"][1] += 1
        STATISTICS["其他"][1] += 1

print(STATISTICS)





# create("../PascalVOC/train/Annotations")
# create("../PascalVOC/test/Annotations")
# create("../PascalVOC/val/Annotations")
